import numpy as np


# 使用梯度下降法求二元二次方程的最小值

# 原函数fx=x02+x12
def function_2(x):
    return x[0] ** 2 + x[1] ** 2


# 求函数梯度值，保存的一维数组向量中
def _numerical_gradient_no_batch(f, x):
    h = 1e-4
    grad = np.zeros_like(x)

    for idx in range(x.size):  # 计算一维向量中每一个元素的梯度值
        tmp_val = x[idx]
        x[idx] = float(tmp_val) + h  # ???
        fxh1 = f(x)  # ?

        x[idx] = float(tmp_val) - h
        fxh2 = f(x)
        grad[idx] = (fxh1 - fxh2) / (2 * h)

        x[idx] = tmp_val
    return grad


def numerical_gradient(f, X):
    if X.ndim == 1:
        return _numerical_gradient_no_batch(f, X)
    else:
        grad = np.zeros_like(X)  # X是二维矩阵，那么从分开求单独的X

        for idx, x in enumerate(X):
            grad[idx] = _numerical_gradient_no_batch(f, X)

    return grad


def gradient_descent(f, init_x, lr=0.01, step_num=100):
    x = init_x
    x_history = []

    for i in range(step_num):
        x_history.append(x.copy())  # 将x一模一样的复制一份

        grad = numerical_gradient(f, x)
        x = x - lr * grad

    return x, np.array(x_history)  # 将一维数组转化为ndarray数据类型


if __name__ == '__main__':
    init_x = np.array([-3.0, 4.0])  # 是这里的问题吗？在上一个实验中真的是这里的问题，数据在定义好回数据类型已经确定好，如果输入浮点数则会四舍五入
    lr = 1e-10
    step_num = 20

    x, x_history = gradient_descent(f=function_2, init_x=init_x, lr=lr, step_num=step_num)

    print('初始化的值是: ' + str(init_x))
    print('x更新的过程如下所示： \n')
    print(x_history)
    print('x下一次更新的值是：')
    print(x)

    # plt.plot([-5, 5], [0, 0], '--b')
    # plt.plot([0, 0], [-5, 5], '--b')
    # plt.plot(x_history[:, 0], x_history[:, 1], 'o')
    #
    # plt.title('gradient_descent Image')
    # plt.xlim(-3.5, 3.5)
    # plt.ylim(-4.5, 4.5)
    # plt.xlabel("X0")
    # plt.ylabel("X1")
    # plt.show()
